#include "yolov5.h"
#include <fstream>
#include <thread>
#include <math.h>
#include <string>
#include <iostream>
#include <iomanip>
#include <chrono>
using namespace std;

#define JSON_ALERT_FLAG_KEY ("is_alert")
#define JSON_ALERT_FLAG_TRUE true
#define JSON_ALERT_FLAG_FALSE false








template <typename T>
static inline bool SortScorePairDescend(const std::pair<float, T> &pair1,
                                        const std::pair<float, T> &pair2)
{
    return pair1.first > pair2.first;
}

#include "timer.hpp"
inline void GetMaxScoreIndex(const std::vector<float> &scores, const float threshold, const int top_k,
                             std::vector<std::pair<float, int>> &score_index_vec)
{
    CV_DbgAssert(score_index_vec.empty());
    // Generate index score pairs.
    for (size_t i = 0; i < scores.size(); ++i)
    {
        if (scores[i] > threshold)
        {
            score_index_vec.push_back(std::make_pair(scores[i], i));
        }
    }

    // Sort the score pair according to the scores in descending order
    std::stable_sort(score_index_vec.begin(), score_index_vec.end(),
                     SortScorePairDescend<int>);

    // Keep top_k scores if needed.
    if (top_k > 0 && top_k < (int)score_index_vec.size())
    {
        score_index_vec.resize(top_k);
    }
}
template <typename T>
static inline float rectOverlap(const T &a, const T &b)
{
    return 1.f - static_cast<float>(jaccardDistance(a, b));
}

template <typename BoxType>
inline void NMSBoxes(const std::vector<BoxType> &bboxes,
                     const std::vector<float> &scores, const float score_threshold,
                     const float nms_threshold, std::vector<int> &indices,
                     float (*computeOverlap)(const BoxType &, const BoxType &),
                     int limit = std::numeric_limits<int>::max())
{
    const float eta = 1.0f;
    const int top_k = 0;
    CV_Assert(bboxes.size() == scores.size());

    // Get top_k scores (with corresponding indices).
    std::vector<std::pair<float, int>> score_index_vec;
    GetMaxScoreIndex(scores, score_threshold, top_k, score_index_vec);

    // Do nms.
    float adaptive_threshold = nms_threshold;
    indices.clear();
    for (size_t i = 0; i < score_index_vec.size(); ++i)
    {
        const int idx = score_index_vec[i].second;
        bool keep = true;
        for (int k = 0; k < (int)indices.size() && keep; ++k)
        {
            const int kept_idx = indices[k];
            float overlap = computeOverlap(bboxes[idx], bboxes[kept_idx]);
            keep = overlap <= adaptive_threshold;
        }
        if (keep)
        {
            indices.push_back(idx);
            if (indices.size() >= limit)
            {
                break;
            }
        }
        if (keep && eta < 1 && adaptive_threshold > 0.5)
        {
            adaptive_threshold *= eta;
        }
    }
}



void datap::reset()
{
    classIds.resize(0);
    confidences.resize(0);
    boxes.resize(0);
    nms_result.resize(0);
}
inline void scale_coords(float *xywh, int newhw, float *sc)
{
    float img1_shape[] = {newhw, newhw};
    float img0_shape[] = {sc[0] * newhw, newhw * sc[1]};
    float gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]);                         //  # gain  = old / new
    float pad[] = {(img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2}; // # wh padding

    float x1 = (xywh[0] - 0.5 * xywh[2]);
    float y1 = (xywh[1] - 0.5 * xywh[3]);

    float x2 = (xywh[0] + 0.5 * xywh[2]);
    float y2 = (xywh[1] + 0.5 * xywh[3]);

    x1 = (x1 - pad[0]) / gain;
    y1 = (y1 - pad[1]) / gain;
    x2 = (x2 - pad[0]) / gain;
    y2 = (y2 - pad[1]) / gain;

    xywh[0] = x1;
    xywh[1] = y1;
    xywh[2] = x2 - x1;
    xywh[3] = y2 - y1;
}
vector<string> getClassNames(const string &imagenet_classes)
{
    ifstream classes_file(imagenet_classes);
    vector<string> classes;
    if (!classes_file.good())
    {
        std::cerr << "ERROR: can't read file with classes names .\n" + imagenet_classes + "\n";
        return classes;
    }
    string class_name;
    while (classes_file.good())
    {
        classes_file >> class_name;
        classes.push_back(class_name);
    }
    classes.pop_back();
    return classes;
}
void datap::press(float *pdata, size_t dlen, int net_width, float conf_thres, float nmsconf, float *sc)
{
    reset();
    for (int i = 0; i < dlen; i++)
    {
        int index = i * net_width;
        if (pdata[index + 4] > conf_thres)
        {
            float *max_cls_pos = std::max_element(pdata + index + 5, pdata + index + net_width);
            float conf = pdata[index + 4] * (*max_cls_pos);
            scale_coords(&pdata[index + 0], 1280, sc);
            confidences.push_back(conf);
            classIds.push_back(int(max_cls_pos - (pdata + index + 5)));
            boxes.push_back(Rect(pdata[index + 0], pdata[index + 1], int(pdata[index + 2]), int(pdata[index + 3])));
            // }
        }
    }

   NMSBoxes<cv::Rect>(boxes, confidences, conf_thres, nmsconf, nms_result,rectOverlap);
}
void datap::print()
{
    std::cout << boxes.size() << " " << nms_result.size() << endl;
    for (auto &it : nms_result)
    {
        std::cout << boxes[it] << " " << confidences[it] << " " << classIds[it] << std::endl;
    }
}
void datap::getbox(vector<Output> &output)
{
    for (int i = 0; i < nms_result.size(); i++)
    {
        int idx = nms_result[i];
        Output result;
        result.id = classIds[idx];
        result.confidence = confidences[idx];
        result.box = boxes[idx];
        
        if (result.box.width<30 ||result.box.height<30)
            continue;
        
        output.push_back(result);
    }
}

Yolov5::Yolov5()
{
    //初始化color
    for (int i = 0; i < className.size(); i++)
    {
        int b = rand() % 256;
        int g = rand() % 256;
        int r = rand() % 256;
        color.push_back(cv::Scalar(b, g, r));
    }
    auto v = getClassNames("/project/train/mysdk/config/class.name");
    v.swap(className);

    string fpath = cv::format("/project/train/mysdk/config/config.yaml");
    // string fpath = cv::format("/project/train/mysdk/config/config.yaml");
    // ACC(fpath.c_str());
    cv::FileStorage fs((fpath), cv::FileStorage::READ);
    // system(string("cat " + fpath).c_str());
    // std::cout << CURRENT_VERSION << std::endl;
    fs["nmsThreshold"] >> nmsThreshold;
    fs["boxThreshold"] >> boxThreshold;

    cout << className.size() << endl;
}
/**
 * @brief Yolov5::postresult_trt 后处理trt的结果
 * @param pdata trt的结果
 * @param img 原生的img
 */

string Yolov5::pack_json(std::vector<Output> &detectedout, bool jsonAlertCode)
{
    Json::Value jRoot;
    Json::Value jAlgoValue;
    Json::Value jDetectValue;

    jAlgoValue["target_info"].resize(0);
    jDetectValue["objects"].resize(0);
    for (auto &obj : detectedout)
    {
        int id_ = int(obj.id);
        auto name = className[id_]; // (obj.id > mConfig.targetRectTextMap[mConfig.language].size() - 1 ? "obj" : mConfig.targetRectTextMap[mConfig.language][obj.id]);
        jDetectValue["objects"].append(obj.toJson(name));
        if (id_ == 0)
        {
            jAlgoValue["target_info"].append(obj.toJson(name));
        }
    }
    jAlgoValue["target_count"] = jAlgoValue["target_info"].size();
    isNeedAlert = jAlgoValue["target_info"].size() > 0;
    jAlgoValue[JSON_ALERT_FLAG_KEY] = isNeedAlert;
    jRoot["algorithm_data"] = jAlgoValue;
    jRoot["model_data"] = jDetectValue;
    Json::StreamWriterBuilder writerBuilder;
    // writerBuilder.settings_["indentation"] = "";
    writerBuilder.settings_["precision"] = 2;
    writerBuilder.settings_["emitUTF8"] = true;
    std::unique_ptr<Json::StreamWriter> jsonWriter(writerBuilder.newStreamWriter());
    std::ostringstream os;
    jsonWriter->write(jRoot, &os);
    mStrOutJson = string(os.str());
    return mStrOutJson;
}

string Yolov5::postresult_trt(float *pdata, cv::Mat img, vector<Output> &output, float thr, int wd)
{
    //求缩放比
    float ratio_h = (float)img.rows / netHeight;
    float ratio_w = (float)img.cols / netWidth;
    float sc[] = {ratio_h, ratio_w};

    cout << boxThreshold << endl;
    dp.press(pdata, 102000, wd, boxThreshold, nmsThreshold, sc);

    dp.getbox(output);
    // cout << "box........." << output.size() << "  " << endl;
    // drawPred(img, output, color);
    static string strout = pack_json(output, false);
    return mStrOutJson;
}

cv::Rect Yolov5::get_rect(cv::Mat &img, float bbox[4])
{
    int l, r, t, b;
    float r_w = netWidth / (img.cols * 1.0);
    float r_h = netHeight / (img.rows * 1.0);
    if (r_h > r_w)
    {
        l = bbox[0] - bbox[2] / 2.f;
        r = bbox[0] + bbox[2] / 2.f;
        t = bbox[1] - bbox[3] / 2.f - (netHeight - r_w * img.rows) / 2;
        b = bbox[1] + bbox[3] / 2.f - (netHeight - r_w * img.rows) / 2;
        l = l / r_w;
        r = r / r_w;
        t = t / r_w;
        b = b / r_w;
    }
    else
    {
        l = bbox[0] - bbox[2] / 2.f - (netWidth - r_h * img.cols) / 2;
        r = bbox[0] + bbox[2] / 2.f - (netWidth - r_h * img.cols) / 2;
        t = bbox[1] - bbox[3] / 2.f;
        b = bbox[1] + bbox[3] / 2.f;
        l = l / r_h;
        r = r / r_h;
        t = t / r_h;
        b = b / r_h;
    }
    return cv::Rect(l, t, r - l, b - t);
}
/**
 * @brief Yolov5::drawPred 画图函数
 * @param img 待检测图片
 * @param result 检测结果
 * @param color 颜色
 */
void Yolov5::drawPred(cv::Mat &img, std::vector<Output> result, std::vector<cv::Scalar> color)
{
    for (int i = 0; i < result.size(); i++)
    {
        int left, top;
        left = result[i].box.x;
        top = result[i].box.y;
        int color_num = i;
        cv::rectangle(img, result[i].box, color[result[i].id], 1, 8);
        string label = className[result[i].id] + ":" + std::to_string(result[i].confidence); //标签
        int baseLine;
        cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.3, 1, &baseLine);
        top = max(top, labelSize.height); //防止超出去
        cv::putText(img, label, cv::Point(left, top), cv::FONT_HERSHEY_SIMPLEX, 1, color[result[i].id], 2);
    }
    // cv::imwrite("D:/Projects/trt.png", img);
    cv::imshow("res", img);
    cv::waitKey(0);
}